from bs4 import BeautifulSoup
from selenium import webdriver
from langchain_community.document_loaders import UnstructuredWordDocumentLoader, PyPDFium2Loader, TextLoader
from langchain_core.documents import Document
import re
import os
import time

def remove_html_comments(html_content):
    comment_pattern = r"<!.*?>"
    html_content = re.sub(comment_pattern, "", html_content, flags=re.DOTALL)
    return html_content.strip()

def remove_css_from_html(html_content):
    # Using BeautifulSoup to parse the HTML content
    soup = BeautifulSoup(html_content, "html.parser")

    [style_tag.decompose() for style_tag in soup.find_all("style")]

    return str(soup)

def get_html_content(url_link):
    chrome_options = webdriver.ChromeOptions()
    chrome_options.add_argument("--headless") # 设置为无头模式
    chrome_options.add_argument("--disable-javascript") # 禁用JavaScript
    driver = webdriver.Chrome(options=chrome_options)
    driver.get(url_link)
    html_content = driver.page_source

    return html_content

# 从html中提取文字
def extract_text_from_link(url_link):
    print("starting extract_text_from_link")
    html_content = get_html_content(url_link)
    html_content = remove_html_comments(html_content)
    html_content_without_css = remove_css_from_html(html_content)
    soup = BeautifulSoup(html_content_without_css, "html.parser")

    texts = [element.strip().replace("\n", " ") for element in soup.find_all(string=True) if element.parent.name !=
             "script" and element.parent.name != "strong" and len(element.strip()) > 0 and element.strip() != "html"]
    texts = ",".join(texts)
    print("end extract_text_from_link")
    return texts

def remove_non_alphanumeric(input_string):
    # 使用正则表达式替换所有非字母和非数字字符为空字符串
    return re.sub(r'[^a-zA-Z0-9\-_]', '', input_string)

def get_text_documents(text, source):
    return Document(page_content=text, metadata={"source": source})

def get_files_documents(file):
    print("starting get_files_documents")
    filename = file.filename
    name, ext = os.path.splitext(filename)

    print(filename)
    file_path = os.path.join("files_cache", name + "_" + str(int(time.time() * 1000)) + ext)
    file.save(file_path)
    if ext == ".pdf":
        loader = PyPDFium2Loader(file_path)
    elif ext == ".txt" or ext == ".md":
        loader = TextLoader(file_path, encoding="utf-8")
    elif ext == ".docx" or ext == ".doc":
        loader = UnstructuredWordDocumentLoader(file_path)

    print("end get_files_documents")
    return loader.load()
